import argparse
import os
import cv2
import torch

from UnetDatasets import UnetDataset
from UnetModel import UnetModel
from config import dataset_path, train_config, project_path


# 预测
def evaluate(epoch, split):
    if epoch == 0:
        epoch = "min_loss"
    test_dataset = UnetDataset(dataset_path=dataset_path, split=split)

    # 定义模型
    model = UnetModel()
    model.to(train_config['device'])

    save_dir = os.path.join(project_path, "result", "weights")
    img_save_dir = os.path.join(project_path, "result", "keypoint_img", split)
    if not os.path.exists(img_save_dir):
        os.makedirs(img_save_dir)

    model_weight_path = os.path.join(save_dir, str(epoch), f"epoch_{str(epoch)}.pkl")
    model.load_state_dict(torch.load(model_weight_path))
    model.eval()

    for index in range(len(test_dataset)):
        img, target_point, img_base, img_name_and_path, window_size = test_dataset.getToVisualization(index=index)
        print(f"index:{index}, file:{img_name_and_path}")
        img = torch.unsqueeze(img, dim=0)  # 训练时采用的是DataLoader函数, 会直接增加第一个维度
        img = img.to(train_config['device'])
        pre = model(img)
        pre = pre.cpu().detach().numpy()
        pre = pre.reshape(pre.shape[0], -1, 2)

        target_point = torch.unsqueeze(target_point, dim=0)
        target_point = target_point.reshape(target_point.shape[0], -1, 2)
        # mask掉没有填充骨头的
        mask = (target_point == 0)
        pre[mask] = 0
        img_with_points = show_point_on_picture(img_base, pre[0], target_point[0], window_size)  # 传入第一个图像的预测结果
        # 显示或保存图像
        img_name = os.path.split(img_name_and_path)[1]

        cv2.imwrite(os.path.join(img_save_dir, img_name.replace('.bmp', '.jpg')), img_with_points)


def show_point_on_picture(img, landmarks, landmarks_gt=None, window_size=None):
    (width, height) = window_size
    # 绘制预测的关键点
    for point in landmarks:
        point = tuple([int(point[0] / 512 * width), int(point[1] / 512 * height)])
        # 没有填充的不显示
        if point == (0, 0):
            continue
        img = cv2.circle(img, center=point, radius=12, color=(0, 0, 255), thickness=-1)

    # 如果有真实标签，则绘制真实标签的关键点
    if landmarks_gt is not None:
        for point in landmarks_gt:
            point = tuple([int(point[0] / 512 * width), int(point[1] / 512 * height)])
            # 没有填充的不显示
            if point == (0, 0):
                continue
            img = cv2.circle(img, center=point, radius=12, color=(0, 255, 0), thickness=-1)

    return img


def run():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epoch', type=int, help='输出调用指定的轮次，默认调用min_loss的', default=0)
    parser.add_argument('--split', type=str, help='使用train还是test进行测试', default="test", choices=["test", "train"])
    # 对一组权重进行预测
    # 解析命令行参数
    args = parser.parse_args()
    print(args)
    evaluate(epoch=args.epoch, split=args.split)


if __name__ == "__main__":
    run()
